package com.tianji.aigc.service.impl;

import cn.hutool.core.date.DateUtil;
import cn.hutool.core.map.MapUtil;
import cn.hutool.core.util.IdUtil;
import cn.hutool.core.util.StrUtil;
import com.tianji.aigc.config.SystemPromptConfig;
import com.tianji.aigc.config.ToolResultHolder;
import com.tianji.aigc.constants.Constant;
import com.tianji.aigc.enums.ChatEventTypeEnum;
import com.tianji.aigc.service.IChatService;
import com.tianji.aigc.service.IChatSessionService;
import com.tianji.aigc.vo.ChatEventVO;
import com.tianji.common.utils.UserContext;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.vectorstore.QuestionAnswerAdvisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
import org.springframework.stereotype.Service;
import reactor.core.publisher.Flux;

import java.util.Map;
import java.util.Optional;
import java.util.concurrent.ConcurrentHashMap;

@Slf4j
@Service
@RequiredArgsConstructor
@ConditionalOnProperty(prefix = "tj.ai", name = "chat-type", havingValue = "ENHANCE")
public class ChatServiceImpl implements IChatService {

    private final ChatClient chatClient;
    private final SystemPromptConfig systemPromptConfig;
    // 作为一个全局的、线程安全的状态存储，用来标记特定 sessionId 对应的对话生成任务是否应该继续进行
    // Key: sessionId，唯一标识一个用户的会话
    // Value: Boolean，true 表示生成任务应该继续，false 或 null (通过 getOrDefault) 表示生成任务应该被停止
    // 目前的版本暂时用Map实现，如果考虑分布式环境的话，可以考虑用redis来实现
    private static final Map<String, Boolean> GENERATE_STATUS = new ConcurrentHashMap<>();
    private final ChatMemory redisChatMemory;
    private final VectorStore vectorStore;
    private final IChatSessionService chatSessionService;

    // 输出结束的标记
    private static final ChatEventVO STOP_EVENT = ChatEventVO.builder().eventType(ChatEventTypeEnum.STOP.getValue()).build();



    /**
     * @param question  问题
     * @param sessionId 会话id
     * @return 聊天结果
     */
    @Override
    public Flux<ChatEventVO> chat(String question, String sessionId) {
        // 获取对话id
        String  conversationId = IChatService.getConversationId(sessionId);
        // 大模型输出内容的缓存器，用于在输出中断后的数据存储
        StringBuilder outputBuilder = new StringBuilder();
        // 获取请求id
        String requestId = IdUtil.fastSimpleUUID();
        // 获取用户id
        Long userId = UserContext.getUser();

//        // 异步更新会话信息
//        chatSessionService.update(sessionId, question, userId);

        // 创建RAG增强
        QuestionAnswerAdvisor questionAnswerAdvisor = QuestionAnswerAdvisor.builder(this.vectorStore)
                .searchRequest(SearchRequest.builder().similarityThreshold(0.6d).topK(6).build())
                .build();

        return chatClient.prompt()
                .system(promptSystemSpec -> promptSystemSpec
                        .text(systemPromptConfig.getChatSystemMessage().get()) // 设置系统提示语
                        .param("now", DateUtil.now())) // 设置当前时间的参数
                .user(question)
                .advisors(advisorSpec -> advisorSpec.advisors(questionAnswerAdvisor) // 创建RAG增强
                        .param(ChatMemory.CONVERSATION_ID, conversationId))
                .toolContext(Map.of(Constant.REQUEST_ID, requestId, Constant.USER_ID, userId))  //通过工具上下文传递参数
                .stream()
                .chatResponse()
                // 以下是新增的用于实现中断的关键操作符
                .doFirst(() -> GENERATE_STATUS.put(sessionId, true)) // 1. 开始时标记
                .doOnComplete(() -> GENERATE_STATUS.remove(sessionId)) //3. 正常完成时清理标记
                .doOnError(throwable -> GENERATE_STATUS.remove(sessionId)) // 2. 出错时清理标记
                .doOnCancel(() -> redisChatMemory.add(conversationId, new AssistantMessage(outputBuilder.toString())))
                // 4.  关键：这是一个过滤和终止操作符，控制Flux流是否继续输出，true：继续，false：终止
                .takeWhile(chatResponse -> GENERATE_STATUS.getOrDefault(sessionId, false))
                .map(chatResponse -> {
                    // 对于响应结果进行处理，如果是最后一条数据，就把此次消息requestId放到内存中
                    // 主要用于存储param（卡片数据）到 redis中，可以根据消息id获取的请求id，再通过请求id就可以获取到参数列表了
                    // 从而解决，在历史聊天记录中没有外参数的问题
                    String finishReason = chatResponse.getResult().getMetadata().getFinishReason();
                    if (StrUtil.equals(Constant.STOP, finishReason)){
                        String messageId = chatResponse.getMetadata().getId();
                        ToolResultHolder.put(messageId, Constant.REQUEST_ID, requestId);
                    }
                    // 获取大模型的输出的内容
                    String text = chatResponse.getResult().getOutput().getText();
                    // 追加到输出内容中
                    outputBuilder.append(text);
                    return ChatEventVO.builder()
                            .eventData(text)
                            .eventType(ChatEventTypeEnum.DATA.getValue())
                            .build();
                })
                .doFinally(signalType -> {
                    // 需要更新对话的标题 或 更新时间
                    var content = StrUtil.format("""
                        ------------
                        USER:{} \n
                        ASSISTANT:{}
                        ------------
                        """, question, outputBuilder.toString());
                    this.chatSessionService.update(sessionId, content, userId);
                })
                .concatWith(// 延迟创建流
                        // 通过请求id获取到参数列表，如果不为空，就将其追加到返回结果中
                        Flux.defer(() -> {
                            Map<String, Object> map = ToolResultHolder.get(requestId);
                            if (MapUtil.isNotEmpty(map)) {
                                // 将参数封装到响应对象中
                                ChatEventVO chatEventVO = ChatEventVO.builder()
                                        .eventType(ChatEventTypeEnum.PARAM.getValue())
                                        .eventData(map)
                                        .build();
                                ToolResultHolder.remove(requestId);
                                return Flux.just(chatEventVO, STOP_EVENT);
                            } else {
                                return Flux.just(STOP_EVENT);
                            }
                        })
                );
    }

    @Override
    public void stop(String sessionId) {
        // 移除标记
        GENERATE_STATUS.remove(sessionId);
    }
}
